package com.tiut.lmap.utils;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.tiut.lmap.common.ErrorCode;
import com.tiut.lmap.exception.ThrowUtils;
import com.tiut.lmap.model.pojo.Chat;
import com.tiut.lmap.service.ChatService;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;

import java.lang.invoke.VarHandle;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

public class InSqlChatMemory implements ChatMemory {
    Map<String, List<Message>> conversationHistory = new ConcurrentHashMap<>();
    @Autowired
    private ChatService chatService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        //将传入的数据转为pojo
        Chat chat = new Chat();
        chat.setId(Long.valueOf(conversationId));
        chat.setType("chat");
        //获取ids集合
        List<Long> ids = chatService.list().stream().map(c -> c.getId()).collect(Collectors.toList());
        //开启新的会话，则删除之前的会话
        if (!ids.contains(chat.getId())) {
            this.conversationHistory.clear();
        }
        this.conversationHistory.putIfAbsent(conversationId, new ArrayList<>());
        this.conversationHistory.get(conversationId).addAll(messages);
        JSONObject json = new JSONObject(conversationHistory);
        chat.setContent(json);
        //如果id存在就更新，反之就保存
        if (ids.contains(chat.getId())) {
            chatService.updateById(chat);
        } else {
            chatService.save(chat);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        Chat cc = chatService.getById(Long.valueOf(conversationId));
        List<Message> all = null;
        if (cc != null) {
            Map<String, Object> map = JSON.parseObject((String) cc.getContent(), Map.class);
            List<JSONObject> jsonMessages = (List<JSONObject>) map.get(conversationId);
            if (jsonMessages != null) {
                all = jsonMessages.stream()
                        .map(json -> {
                           String text = String.valueOf(json.get("text"));
                            String role = String.valueOf(json.get("messageType"));
                            if ("USER".equals(role)) {
                                return new UserMessage(text);
                            } else if ("ASSISTANT".equals(role)) {
                                return new AssistantMessage(text);
                            }
                            return null; // 其他情况可以扩展
                        })
                        .filter(Objects::nonNull)
                        .collect(Collectors.toList());
            }
        }
//        List<Message> all = this.conversationHistory.get(conversationId);
        return all != null ? all.stream().skip(Math.max(0, all.size() - lastN)).toList() : List.of();
    }

    @Override
    public void clear(String conversationId) {
        this.conversationHistory.remove(conversationId);
    }
}
